Greptile Summary

This PR enables torch.compile(fullgraph=True) for FP8 autocast by making two structural changes to transformer_engine/pytorch/quantization.py: (1) moving all mutable FP8 global state off class attributes into a FP8GlobalState dataclass singleton (quantization_state), so writes no longer hit setattr(type, ...) which torch.compile cannot trace; and (2) replacing functools.lru_cache-wrapped support-check functions with explicit module-level caches guarded by @torch.compiler.assume_constant_result, preventing the compiler from tracing into hardware-capability checks. DelayedScaling is explicitly blocked under torch.compile with a clear error. The remaining 9 files are mechanical migrations from the old class-attribute spellings to the new quantization_state.* accessor.

Key issues identified (not yet discussed in prior threads):

• Private PyTorch internal API imports in test_torch_compile.py (lines 10-15): torch._opaque_base.OpaqueBaseMeta and torch._library.opaque_object.* are underscore-prefixed internals with no stability contract. A future PyTorch minor release could move or rename them, silently breaking every test in the file with an ImportError. The imports should be guarded with a try/except and a version-gating comment, or replaced by any equivalent public surface when available.

• test_autocast_sanity uses torch.nn.Linear instead of ToyLinear (lines 310-329): Wrapping a plain torch.nn.Linear in te.autocast does not trigger any TE quantization. All four recipe variants (Float8CurrentScaling, Float8BlockScaling, MXFP8BlockScaling, NVFP4BlockScaling) execute identical BF16 code, so the test provides zero evidence that any recipe actually works end-to-end under torch.compile. The ToyLinear module defined in the same file (precisely for this purpose) should be used instead.

Previously-flagged issues still open: get_default_fp8_recipe() called unconditionally in autocast_enter when recipe=None (will assert under torch.compile); reset() swapping the singleton object invalidating compiled-graph guards; assert in the compile branch of fp8_graph_capturing stripped by -O; id()-based autocast keys risking unbounded dict growth; check_recipe_support not guarding the recipe=None + enabled=True + compile case; and get_fp8_recipe() falling through to get_default_fp8_recipe() when called without an active autocast context during compilation.

Confidence Score: 2/5

• Several correctness issues in the compile path remain unaddressed, and the new test suite provides weaker coverage than it appears to.
• Score reflects two categories of open risk: (1) The autocast_enter function unconditionally calls get_default_fp8_recipe() when recipe=None, which asserts not torch.compiler.is_compiling(). This means calling te.autocast() (no recipe) inside torch.compile — even with enabled=False — raises an opaque AssertionError rather than a user-friendly error. (2) FP8GlobalStateManager.reset() replaces the quantization_state singleton with a new object, invalidating all torch.compile guards that captured the old object's identity, causing full recompilation on every subsequent call. Additionally, the two new test issues identified in this review mean the coverage picture for "recipes work under torch.compile" is not as complete as the PR description implies. The mechanical migrations across the other 8 files look correct.
• transformer_engine/pytorch/quantization.py (lines 317, 469, 486, 574, 591) and tests/pytorch/test_torch_compile.py (lines 10–15, 310–329) need the most attention before merging.

Important Files Changed

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["te.autocast(recipe=R, enabled=E)"] --> B{enabled?}
    B -- yes --> C["check_recipe_support(recipe)"]
    C --> D{is_compiling AND\nDelayedScaling?}
    D -- yes --> E["RuntimeError ✗"]
    D -- no --> F["save fp8_state tuple\nfrom quantization_state.*"]
    B -- no --> F
    F --> G["autocast_enter(enabled, recipe, ...)"]
    G --> H{recipe is None?}
    H -- yes --> I["get_default_fp8_recipe()\nassert not is_compiling ⚠️"]
    H -- no --> J["quantization_state.autocast_arguments[key] = recipe,group\nquantization_state.fp8_* = new values"]
    I --> J
    J --> K["yield — user code runs"]
    K --> L["finally: restore quantization_state.*\nfrom saved fp8_state tuple"]
    L --> M["autocast_exit: depth--, reduce amaxes if needed"]

    subgraph "support cache (module-level globals)"
        N["_FP8_SUPPORT / _MXFP8_SUPPORT\n_NVFP4_SUPPORT / _FP8_BLOCK_SCALING_SUPPORT"]
        O["check_fp8_support() etc.\n@assume_constant_result"] --> N
    end

    subgraph "FP8GlobalState singleton"
        P["quantization_state\n.fp8_enabled / .fp8_recipe\n.autocast_depth / .is_first_fp8_module\n.global_amax_buffer / .autocast_arguments\n..."]
    end

    J -.-> P
    L -.-> P

_{Reviews (5): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

/te-ci pytorch L1

/te-ci pytorch